import cv2
import numpy as np
import os

# 全局变量，用于保存鼠标点击的像素坐标
clicked_point = None

def mouse_callback(event, x, y, flags, param):
    """鼠标点击回调函数"""
    global clicked_point
    if event == cv2.EVENT_LBUTTONDOWN:  # 左键点击
        clicked_point = (x, y)
        print(f"选中的像素点坐标: {clicked_point}")

def get_xyz_from_pixel(pixel: tuple, depth_map: np.ndarray, K: np.ndarray, D: np.ndarray) -> tuple:
    """根据像素点坐标获取点云空间内的 (x, y, z)"""

    # 获取的是 x y
    u, v = pixel  # x,y
    print(f'ori {u, v}')
    # uv 行 列
    # u = int( u / 768 * 480)
    # v = int( v / 1024 * 640)
    print(f'cvt {u, v}')


    # 获取深度值（单位：米）
    # z = depth_map[u, v] / 1000.0  # 假设深度图单位是毫米，转换为米
    z = depth_map[v][u] / 1000.0  # 假设深度图单位是毫米，转换为米
    print(f'z {z}')
    if z <= 0 or z > 20.0:  # 检查深度值是否有效
    

        neighbors = [
            (u - 1, v - 1), (u, v - 1), (u + 1, v - 1),  # 上方三个点
            (u - 1, v),                 (u + 1, v),      # 左右两个点
            (u - 1, v + 1), (u, v + 1), (u + 1, v + 1)   # 下方三个点
        ]
        
        for nu, nv in neighbors:
            if 0 <= nv < depth_map.shape[0] and 0 <= nu < depth_map.shape[1]:  # 检查是否在图像范围内
                neighbor_z = depth_map[nv][nu] / 1000.0  # 获取深度值
                print(f"邻居像素点 ({nu}, {nv}) 的深度值: {neighbor_z} 米")
            else:
                print(f"邻居像素点 ({nu}, {nv}) 超出图像范围")
                      
        raise ValueError(f"无效的深度值: {z} m")
        
    
    # 使用鱼眼模型进行反投影
    pts = np.array([[[u, v]]], dtype=np.float32)  # 像素点坐标需要是 (1, 1, 2) 的形状
    undistorted_pts = cv2.fisheye.undistortPoints(pts, K, D)  # 去畸变并反投影
    x = undistorted_pts[0][0][0] * z
    y = undistorted_pts[0][0][1] * z
    return x, y, z

def load_depth_map(depth_path: str, depth_height: int = 480, depth_width: int = 640) -> np.ndarray:
    with open(depth_path, 'rb') as f:
        depth_data = np.fromfile(f, dtype=np.uint16)

    depth_map = depth_data.reshape((depth_height, depth_width))
    return depth_map


def compute_depth_pseudo_color(depth_image: np.ndarray, max_depth: float = 5.0) -> np.ndarray:
    """将深度图转为伪彩色图像"""
    depth_normalized = np.clip(depth_image / (max_depth * 1000), 0, 1)
    depth_colored = cv2.applyColorMap((depth_normalized * 255).astype(np.uint8), cv2.COLORMAP_JET)
    return depth_colored

def main():
    global clicked_point  # 声明 clicked_point 为全局变量

    K = np.array([[380.514984, 0, 512.130066],
                  [0, 380.059082, 332.302795],
                  [0, 0, 1]]) 
    # 380.977051, 380.876495, 518.386719, 337.40976
    # D = np.array([0.045951847, -0.011027799, 0.007337844, -0.00242689601])
    D = np.array([0.0422173887, -0.000108877211, -0.000979754375, -0.000220517832])


    # 文件路径（用户可修改为自己的路径）
    # depth_path = input("请输入深度图路径（.raw 文件）：").strip()
    # rgb_path = input("请输入 RGB 图像路径：").strip()
    # /media/ai/0bea6433-71ce-4bf1-a689-3b0348c3c57b/vslam/0423/20250409080003_92965/depth/1744185603703480.raw
    # /media/ai/0bea6433-71ce-4bf1-a689-3b0348c3c57b/vslam/0423/20250409080003_9
    depth_path = '/media/ai/0bea6433-71ce-4bf1-a689-3b0348c3c57b/vslam/0423/20250409080003_92965/depth/1744185603703480.raw'
    depth_path = '/home/ai/wlm/gitee/map_learing/01_binocular_vision_xuce/data/1744185603703480.raw'
    # rgb_path = '/media/ai/0bea6433-71ce-4bf1-a689-3b0348c3c57b/vslam/0423/20250409080003_92965/results/left_fisheye/1744185603703480.jpg'
    rgb_path = '/home/ai/wlm/gitee/map_learing/01_binocular_vision_xuce/pseudo1024x768.jpg'
    rgb_path = '/home/ai/wlm/gitee/map_learing/01_binocular_vision_xuce/data/contour_0.png'
    # depth_path = '/home/ai/wlm/gitee/map_learing/01_binocular_vision_xuce/data/1744185603703480.raw'
    # rgb_path = '/home/ai/wlm/gitee/map_learing/01_binocular_vision_xuce/data/1744185603703480.jpg'


    depth_path = '/home/ai/wlm/gitee/map_learing/01_binocular_vision_xuce/imgs_raw/1744185747802609.raw'
    rgb_path = '/home/ai/wlm/gitee/map_learing/01_binocular_vision_xuce/imgs_rep/1744185747802609.jpg'


    # 加载深度图
    depth_map = load_depth_map(depth_path)
    if depth_map is None:
        print("深度图加载失败，程序退出。")
        return
    

    # 打开一个txt文件准备写入
    with open("depth_map.txt", "w") as file:
        # 遍历二维数组
        for i in range(depth_map.shape[0]):  # 遍历行
            for j in range(depth_map.shape[1]):  # 遍历列
                # 获取坐标和对应的值
                coord = (i, j)
                value = depth_map[i, j]
                # 打印到控制台
                # print(f"坐标: {coord}, 值: {value}")
                if value != 0:
                # 写入到文件
                    file.write(f"坐标: {coord}, 值: {value/1000.0}\n")

    print("数据已成功写入到 depth_map.txt 文件中。")
    # save
    pseudo_color_depth = compute_depth_pseudo_color(depth_map)
    cv2.imwrite("pseudo.jpg",pseudo_color_depth )


    resized_depth_map = cv2.resize(depth_map, (1024, 768), interpolation=cv2.INTER_NEAREST)


    pseudo_color_depth1 = compute_depth_pseudo_color(resized_depth_map)
    cv2.imwrite("pseudo1.jpg",pseudo_color_depth1 )

    # 打开一个txt文件准备写入
    with open("depth_map_change.txt", "w") as file:
        # 遍历二维数组
        for i in range(resized_depth_map.shape[0]):  # 遍历行
            for j in range(resized_depth_map.shape[1]):  # 遍历列
                # 获取坐标和对应的值
                coord = (i, j)
                value = resized_depth_map[i, j]
                # 打印到控制台
                # print(f"坐标: {coord}, 值: {value}")
                if value != 0:
                # 写入到文件
                    file.write(f"坐标: {coord}, 值: {value/1000.0}\n")

    print("数据已成功写入到 depth_map_change.txt 文件中。")
    # save
    pseudo_color_depth = compute_depth_pseudo_color(depth_map)
    cv2.imwrite("pseudo.jpg",pseudo_color_depth )


    # 加载 RGB 图像
    if not os.path.exists(rgb_path):
        print("RGB 图像路径无效，程序退出。")
        return
    rgb_image = cv2.imread(rgb_path)
    if rgb_image is None:
        print("无法加载 RGB 图像，程序退出。")
        return

    # 确保 RGB 图像和深度图尺寸一致
    depth_height, depth_width = depth_map.shape
    # rgb_image = cv2.resize(rgb_image, (depth_width, depth_height))

    # 显示 RGB 图像并绑定鼠标回调
    cv2.namedWindow("RGB Image")
    cv2.setMouseCallback("RGB Image", mouse_callback)

    while True:
        # 显示 RGB 图像
        display_image = rgb_image.copy()

        if clicked_point is not None:
            try:
                # 获取点击点的 (x, y, z)
                x, y, z = get_xyz_from_pixel(clicked_point, resized_depth_map, K, D)
                # x, y, z = get_xyz_from_pixel(clicked_point, resized_depth_map, K, D)
                print(f"像素点 {clicked_point} 对应的点云坐标: x={x:.3f}, y={y:.3f}, z={z:.3f}")
                
                # 在图像上显示深度值和点云坐标
                cv2.putText(display_image, f"Depth: {z:.3f}m", (clicked_point[0], clicked_point[1] - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
                cv2.putText(display_image, f"X: {x:.3f}, Y: {y:.3f}, Z: {z:.3f}", (10, 30),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                clicked_point = None  # 重置点击点
            except ValueError as e:
                print(e)
                clicked_point = None  # 重置点击点

        cv2.imshow("RGB Image", display_image)
        key = cv2.waitKey(1) & 0xFF

        if key == 27:  # 按下 ESC 键退出
            break

    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
